import argparse


def get_config():
    parser = argparse.ArgumentParser(
        description='onpolicy', formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("--seed", type=int, default=1, help="numpy/torch的随机种子")
    parser.add_argument("--cuda", action='store_false', default=True,
                        help="默认True，使用GPU训练；否则使用CPU")
    parser.add_argument("--cuda_deterministic",
                        action='store_false', default=True,
                        help="默认确保随机种子有效。如果设置，则绕过此功能")
    parser.add_argument("--n_training_threads", type=int,
                        default=1, help="训练时torch线程数量")
    parser.add_argument("--n_rollout_threads", type=int, default=1,
                        help="训练时并行环境数量")
    parser.add_argument("--n_eval_rollout_threads", type=int, default=1,
                        help="评估时并行环境数量")
    parser.add_argument("--num_env_steps", type=int, default=10e6,
                        help='训练的环境步数（默认：10e6）')

    # 经验回放缓冲区参数
    parser.add_argument("--episode_length", type=int,
                        default=200, help="每个episode的最大长度")

    # 优化器参数
    parser.add_argument("--lr", type=float, default=5e-4,
                        help='学习率（默认：5e-4）')
    parser.add_argument("--critic_lr", type=float, default=5e-4,
                        help='评论家网络学习率（默认：5e-4）')
    parser.add_argument("--opti_eps", type=float, default=1e-5,
                        help='RMSprop优化器epsilon（默认：1e-5）')
    parser.add_argument("--weight_decay", type=float, default=0,
                        help='权重衰减（默认：0）')

    # PPO参数
    parser.add_argument("--ppo_epoch", type=int, default=15,
                        help='PPO训练轮数（默认：15）')
    parser.add_argument("--clip_param", type=float, default=0.2,
                        help='PPO裁剪参数（默认：0.2）')
    parser.add_argument("--num_mini_batch", type=int, default=1,
                        help='PPO小批次数量（默认：1）')
    parser.add_argument("--entropy_coef", type=float, default=0.01,
                        help='熵项系数（默认：0.01）')
    parser.add_argument("--value_loss_coef", type=float,
                        default=1, help='价值损失系数（默认：1）')
    parser.add_argument("--max_grad_norm", type=float, default=10.0,
                        help='梯度最大范数（默认：10.0）')
    parser.add_argument("--gamma", type=float, default=0.99,
                        help='奖励折扣因子（默认：0.99）')
    parser.add_argument("--gae_lambda", type=float, default=0.95,
                        help='GAE lambda参数（默认：0.95）')
    parser.add_argument("--huber_delta", type=float, default=10.0, help="Huber损失系数")

    # 保存参数
    parser.add_argument("--save_interval", type=int, default=100,
                        help="连续两次模型保存之间的时间间隔")

    # 日志参数
    parser.add_argument("--log_interval", type=int, default=5,
                        help="连续两次日志打印之间的时间间隔")

    # 评估参数
    parser.add_argument("--use_eval", action='store_true', default=False,
                        help="默认不启动评估。如果设置，则与训练同时启动评估")
    parser.add_argument("--eval_interval", type=int, default=25,
                        help="连续两次评估进度之间的时间间隔")
    parser.add_argument("--eval_episodes", type=int, default=32, help="单次评估的episode数量")

    # 预训练参数
    parser.add_argument("--model_dir", type=str, default=None,
                        help="默认None。设置预训练模型的路径")

    # Transformer相关参数
    parser.add_argument("--n_block", type=int, default=1,
                        help="Transformer块数量")
    parser.add_argument("--n_embd", type=int, default=64,
                        help="嵌入维度")
    parser.add_argument("--n_head", type=int, default=1,
                        help="注意力头数量")
    
    # 网络配置参数
    parser.add_argument('--dropout_rate', type=float, default=0.1,
                        help='Dropout比率')

    # Map相关参数（经度、纬度）
    parser.add_argument("--map_lon_min", type=float, default=73.0,
                        help="地图最小经度（默认：73.0）")
    parser.add_argument("--map_lon_max", type=float, default=135.0,
                        help="地图最大经度（默认：135.0）")
    parser.add_argument("--map_lat_min", type=float, default=18.0,
                        help="地图最小纬度（默认：18.0）")
    parser.add_argument("--map_lat_max", type=float, default=54.0,
                        help="地图最大纬度（默认：54.0）")

    # 威胁区域参数（半球形威胁区）
    parser.add_argument("--threat_center_lon", type=float, default=119.0,
                        help="威胁区中心经度（默认：119.0）")
    parser.add_argument("--threat_center_lat", type=float, default=24.5,
                        help="威胁区中心纬度（默认：24.5）")
    parser.add_argument("--threat_center_alt", type=float, default=0.0,
                        help="威胁区中心高度（米，默认：0.0）")
    parser.add_argument("--threat_damage_radius", type=float, default=50000.0,
                        help="威胁区毁伤有效半径（米，默认：50000.0）")

    # 海拔高度参数
    parser.add_argument("--altitude_min", type=float, default=0.0,
                        help="最低海拔高度（米，默认：0.0）")
    parser.add_argument("--altitude_max", type=float, default=100000.0,
                        help="最高海拔高度（米，默认：100000.0）")

    return parser
